# -*- coding: utf-8 -*-
'''
Created on 20.08.2019

@author: yu03
'''
import datetime
from PyUeye import file_name, file_path, txt_name, doc_name
import numpy as np
import os
import matplotlib.pyplot as plt
from scipy import signal
from FFT_Interpolation import FFT_interpolation_boxcar, FFT_interpolation_2, FFT_interpolation_compare, FFT_cal
from scipy.optimize import curve_fit
from scipy import signal

def Export_Data(file_name, header, out_str):
    print('Writing Data')
    with open(file_name,'w') as fid: ######################################################################################
        fid.writelines(header)
        fid.writelines(out_str)
    print('Finish Writing')
    return

def running_mean(x, N):
    cumsum = np.cumsum(np.insert(x, 0, 0)) 
    return np.concatenate(( np.zeros(int((N)/2)), (cumsum[N:]-cumsum[:-N])/float(N) , np.zeros(int((N-1)/2)) ))
def fit_func(x, a, b, c):
    return a*(x-b)**2 + c
# p= -0.0213125
p= -0.02135
def fit_func_phase(x, q):
    return p*x+q

def fit_func_line(x, p, q):
    return p*x+q

now = datetime.datetime.now()
#############
# fs_cam = 575.8 ###USB2.0
# fs_cam = 1370.72 ###USB3.0
fs_cam = 1000 ### Hardware Trigger

Lamda = 633e-9
pix_size = 5.3e-6
V_x, V_y, V_z = 0, 0, 0

img_set = []
hor_phase_centers = []
hor_freq_set = []
hor_f_fit_set = []
hor_phi_fit_set = []


f = open(file_name, 'rb')
f_size = os.fstat(f.fileno()).st_size
frames = []
time_sequence = []
while f.tell() < f_size:
#         print(f.tell())
    frame = np.load(f, allow_pickle=True)
    frames.append(frame[0])
    time_sequence.append(frame[1])

# x_axis = np.linspace(0, len(time_sequence)-1, len(time_sequence))

frames = np.array(frames)
print(frames.shape)

print('file_size: %.1f MB'%(f_size/1e6))
print('%i frames'%len(frames))
print('%i lines/frame'%len(frames[0]))
print('%i pixels/line'%len(frames[0][0]))
print('%i channels/pixel'%len(frames[0][0][0]))
print('frame rate: %.2f'%fs_cam)
print('record time: %f s'%(len(frames)/fs_cam))

num = 0
for frame in frames[::]:
    num += 1
    hor_phases = []
    hor_freqs = []
    for line in frame:
        line = np.array([i[0]+256*i[1] for i in line])
        
#         line = np.diff(line.astype('float32'))
#         line = np.concatenate(([0], line))

        hor_center = line
        DC_num = 1000
        hor_freq_estim, hor_phase_estim, hor_freqline, hor_sig_magnitude, hor_sig_phase,  hor_m_k_num, hor_X_m_k, hor_freq_for_phase = FFT_interpolation_2(hor_center, pix_size, 1e5, DC_num)
        hor_FFT_start = np.where(hor_sig_magnitude[DC_num:] > hor_X_m_k*0.4)[0][0]+DC_num
        hor_FFT_end = np.where(hor_sig_magnitude[DC_num:] > hor_X_m_k*0.4)[0][-1]+DC_num
        hor_fit_x = hor_freqline[hor_FFT_start:hor_FFT_end+1]
        hor_fit_y = hor_sig_magnitude[hor_FFT_start:hor_FFT_end+1]
        hor_fit_phase = hor_sig_phase[hor_FFT_start:hor_FFT_end+1]
        hor_params = curve_fit(fit_func, hor_fit_x, hor_fit_y)
        [hor_a, hor_b, hor_c] = hor_params[0]
        hor_f_fit = hor_b
        hor_phases.append(hor_phase_estim)
        hor_freqs.append(hor_f_fit)
    hor_phases = np.array(hor_phases)
    hor_freqs = np.array(hor_freqs)
    hor_f_fit_set.append(np.sum(hor_freqs)/16)
    fitted_phase = np.sum(np.unwrap(hor_phases))/16
    if fitted_phase > np.pi:
        fitted_phase = fitted_phase - 2*np.pi
    if fitted_phase < -np.pi:
        fitted_phase = fitted_phase + 2*np.pi
    hor_phase_centers.append(fitted_phase)
    print(num, hor_m_k_num)

hor_f_fit_set = np.array(hor_f_fit_set)
hor_phase_centers = np.array(hor_phase_centers)
hor_angle = (V_x-Lamda*hor_f_fit_set/2)*1e6 ### urad
hor_length = np.unwrap(hor_phase_centers)/4/np.pi*Lamda*1e9 ### nm

''' 
    输出数据
'''

header = ['%s\n' %doc_name,
      'Local current time : %s\n' %now.strftime("%Y-%m-%d %H:%M:%S"),
      'Fs = %e (Hz)\n' %fs_cam,##########################################################################################################
      'Record time: %e (s)\n' %(len(frames)/fs_cam),############################################################################################
      'Frame Number = %i\n' %len(frames),############################################################################################
      'Lines/Frame = %i\n' %len(frames[0]),############################################################################################
      'Pixels/Line = %i\n' %len(frames[0][0]),############################################################################################
      'Channels/Pixel = %i\n' %len(frames[0][0][0]),############################################################################################
      'Channel_1: Horizontal freq.(16-line averaged)\n',############################################################################################
      'Channel_2: Horizontal phase(16-line averaged)\n',############################################################################################
      'Channel_3: XXX\n',############################################################################################
      'Channel_4: XXX\n',############################################################################################
#       'Channel_3: Demodulated S_curve (from CH_No.2)\n',############################################################################################
#       'Channel_4: Demodulated S_curve (from CH_No.1)\n',############################################################################################
#       'Data_Ch5: Difference between Channel_3 & Channel_4 (Ch3-Ch4)\n'################################################################################
      '-------------------------------------------------\n',
      ]
# out_str = ['%.4f, %.4f, %.4f, %.4f\n' %(Data_Ch1[i], Data_Ch2[i], Data_Ch3[i], Data_Ch4[i]) for i in range(len(Data))] 
# out_str = [' %.4f, %.4f, %.4f, %.4f, %.4f\n' %(Data_Ch1[i], Data_Ch2[i], Data_Ch3[i], Data_Ch4[i], Data_Ch5[i]) for i in range(len(Data))]    
out_str = ['%f, %f\n' %(hor_f_fit_set[i], hor_phase_centers[i]) for i in range(len(hor_f_fit_set))]    
 
''' 
    保存文件
'''
Export_Data(txt_name, header, out_str)
print('TXT file saved')


if 1: ### Plotting Noise
    
   Noise_PSD########### linearity range
    start_cutting, stop_cutting = 0, len(hor_phase_centers)
#     start_cutting, stop_cutting = 20, 3000
    
    selected_range = np.linspace(start_cutting, stop_cutting, stop_cutting-start_cutting)
    
    ############### Plotting
     
    plt.figure(1)
    plt.gcf().set_size_inches(18,9)
    ax1 = plt.subplot(3,2,1)
    plt.plot(hor_angle, color='blue', marker=' ')
    plt.title("Horizontal Tilting")
    plt.ylabel("Horizontal Angle (urad)")
    plt.xlabel("Samples")
    plt.grid(which='both', axis='both')
        
    ax2 = plt.subplot(3,2,2)
    plt.plot(hor_length, color='blue', marker=' ')
    plt.title("Horizontal Phase")
    plt.ylabel("Horizontal Length (nm)]")
    plt.xlabel("Samples")
    plt.grid(which='both', axis='both')
      
    ax3 = plt.subplot(3,2,3)
    plt.plot(hor_angle, color='blue', marker=' ')
    plt.title("Horizontal Tilting")
    plt.ylabel("Horizontal Angle (urad)")
    plt.xlabel("Samples")
    plt.grid(which='both', axis='both')
       
    ax4 = plt.subplot(3,2,4)
    plt.plot(hor_length, color='blue', marker=' ')
    plt.title("Horizontal Length")
    plt.ylabel("Horizontal Length (nm)")
    plt.xlabel("Samples")
    plt.grid(which='both', axis='both')

    ax5 = plt.subplot(3,2,5)
    f_line, PS = signal.welch(hor_angle, fs_cam, nperseg=len(hor_angle))
    PSD = np.sqrt(PS)
    plt.loglog(f_line, PSD, 'b')
    plt.title("Horizontal Tilting Noise")
    pltNoise_PSDel("PSD (urad/$\sqrt{Hz}$)")
    plt.xlabel("Freq. (Hz)")
    plt.grid(which='both', axis='both')
    plt.ylim(1e-6, 10) # for Ramp/20s noise
#     plt.ylim(1e-4, 0.1) # for noise
    
    ax6 = plt.subplot(3,2,6)
    f_line, PS = signal.welch(hor_length, fs_cam, nperseg=len(hor_length))
    PSD = np.sqrt(PS)
    plt.loglog(f_line, PSD, 'b')
    plt.title("Horizontal Length Noise")
    pltNoise_PSDel("PSD (nm/$\sqrt{Hz}$)")
    plt.xlabel("Freq. (Hz)")
    plt.grid(which='both', axis='both')
    plt.ylim(1e-6, 10) # for Ramp/20s noise
#     plt.ylim(1e-6, 2e-2) # for square wave
#     plt.ylim(1e-6, 1e-3)# for noise
    
    plt.tight_layout()
    plt.show()



if 0: ### Plotting Nonlinearity
    
    ############### linearity range
    start_cutting, stop_cutting = 0, len(hor_phase_centers)
#     start_cutting, stop_cutting = 20, 3000
    
    selected_range = np.linspace(start_cutting, stop_cutting, stop_cutting-start_cutting)
    
    ############### Plotting
     
    plt.figure(1)
    plt.gcf().set_size_inches(18,9)
    ax1 = plt.subplot(4,2,1)
    plt.plot(hor_angle, color='blue', marker=' ')
    measured_angle = hor_angle[start_cutting:stop_cutting]
    [p, q] = np.polyfit(np.linspace(0, len(measured_angle)-1, len(measured_angle)), measured_angle, 1)
    linear_fit = np.linspace(0, len(measured_angle)-1, len(measured_angle))*p + q
    hor_angle_nonlinearity = (measured_angle - linear_fit)
    plt.plot(selected_range, linear_fit, color='red')
    plt.title("Horizontal Tilting")
    plt.ylabel("Horizontal Angle (urad)")
    plt.xlabel("Samples")
    plt.grid(which='both', axis='both')
        
    ax2 = plt.subplot(4,2,2)
    plt.plot(hor_length, color='blue', marker=' ')
    measured_length = hor_length[start_cutting:stop_cutting]
    [p, q] = np.polyfit(np.linspace(0, len(measured_length)-1, len(measured_length)), measured_length, 1)
    linear_fit = np.linspace(0, len(measured_length)-1, len(measured_length))*p + q
    hor_length_nonlinearity = (measured_length - linear_fit)
    plt.plot(selected_range, linear_fit, color='red')
    plt.title("Horizontal Phase")
    plt.ylabel("Horizontal Length (nm)")
    plt.xlabel("Samples")
    plt.grid(which='both', axis='both')
      
    ax3 = plt.subplot(4,2,3)
    plt.plot(selected_range, hor_angle_nonlinearity, color='blue')
    plt.title("Horizontal Tilting Nonlinearity")
    plt.ylabel("Horizontal Angle (urad)")
    plt.xlabel("Samples")
    plt.grid(which='both', axis='both')
       
    ax4 = plt.subplot(4,2,4)
    plt.plot(selected_range, hor_length_nonlinearity, color='blue')
    plt.title("Horizontal Length Nonlinearity")
    plt.ylabel("Horizontal Length (nm)")
    plt.xlabel("Samples")
    plt.grid(which='both', axis='both')
      
    ax5 = plt.subplot(4,2,5)
    plt.plot(selected_range, hor_angle_nonlinearity, color='blue')
    plt.title("Horizontal Tilting Nonlinearity")
    plt.ylabel("Horizontal Angle (urad)")
    plt.xlabel("Samples")
    plt.grid(which='both', axis='both')
      
    ax6 = plt.subplot(4,2,6)
    plt.plot(selected_range, hor_length_nonlinearity, color='blue')
    plt.title("Horizontal Length Nonlinearity")
    plt.ylabel("Horizontal Length (nm)")
    plt.xlabel("Samples")
    plt.grid(which='both', axis='both')
      
    ax1 = plt.subplot(4,2,7)
    f_line, PS = signal.welch(hor_angle_nonlinearity, fs_cam, nperseg=len(hor_angle_nonlinearity))
    PSD = np.sqrt(PS)
    plt.loglog(f_line, PSD, 'b')
    plt.title("Horizontal Tilting Noise")
    pltNoise_PSDel("PSD (urad/r'$\sqrt{Hz}$'")
    plt.xlabel("Freq. (Hz)")
    plt.grid(which='both', axis='both')
    plt.ylim(1e-6, 10) # for Ramp/20s noise
#     plt.ylim(1e-4, 0.1) # for noise
    
    ax2 = plt.subplot(4,2,8)
    f_line, PS = signal.welch(hor_length_nonlinearity, fs_cam, nperseg=len(hor_length_nonlinearity))
    PSD = np.sqrt(PS)
    plt.loglog(f_line, PSD, 'b')
    plt.title("Horizontal Length Noise")
    pltNoise_PSDel("PSD (nm/sqrt(Hz))")
    plt.xlabel("Freq. (Hz)")
    plt.grid(which='both', axis='both')
    plt.ylim(1e-6, 10) # for Ramp/20s noise
#     plt.ylim(1e-6, 2e-2) # for square wave
#     plt.ylim(1e-6, 1e-3)# for noise
    
    plt.tight_layout()
    plt.show()



